"""
图表组件
提供各种数据可视化图表组件
"""

import streamlit as st
import plotly.graph_objects as go
import plotly.express as px
import pandas as pd
from typing import Dict, List, Any, Optional, Union
import numpy as np
from datetime import datetime, timedelta


def create_line_chart(
    data: Union[pd.DataFrame, List[Dict]],
    x_col: str,
    y_col: str,
    title: str = "",
    color: str = "#1f77b4",
    height: int = 400,
    show_points: bool = False
) -> go.Figure:
    """
    创建折线图
    
    Args:
        data: 数据
        x_col: X轴列名
        y_col: Y轴列名
        title: 图表标题
        color: 线条颜色
        height: 图表高度
        show_points: 是否显示数据点
    
    Returns:
        Plotly图表对象
    """
    if isinstance(data, list):
        df = pd.DataFrame(data)
    else:
        df = data
    
    fig = go.Figure()
    
    mode = "lines+markers" if show_points else "lines"
    
    fig.add_trace(go.Scatter(
        x=df[x_col],
        y=df[y_col],
        mode=mode,
        line=dict(color=color, width=2),
        marker=dict(size=6, color=color),
        name=y_col
    ))
    
    fig.update_layout(
        title=title,
        xaxis_title=x_col,
        yaxis_title=y_col,
        height=height,
        showlegend=False,
        margin=dict(l=50, r=50, t=50, b=50),
        plot_bgcolor='rgba(0,0,0,0)',
        paper_bgcolor='rgba(0,0,0,0)'
    )
    
    fig.update_xaxes(showgrid=True, gridwidth=1, gridcolor='lightgray')
    fig.update_yaxes(showgrid=True, gridwidth=1, gridcolor='lightgray')
    
    return fig


def create_multi_line_chart(
    data: Union[pd.DataFrame, List[Dict]],
    x_col: str,
    y_cols: List[str],
    title: str = "",
    colors: Optional[List[str]] = None,
    height: int = 400
) -> go.Figure:
    """
    创建多线图
    
    Args:
        data: 数据
        x_col: X轴列名
        y_cols: Y轴列名列表
        title: 图表标题
        colors: 颜色列表
        height: 图表高度
    
    Returns:
        Plotly图表对象
    """
    if isinstance(data, list):
        df = pd.DataFrame(data)
    else:
        df = data
    
    if colors is None:
        colors = px.colors.qualitative.Set1[:len(y_cols)]
    
    fig = go.Figure()
    
    for i, y_col in enumerate(y_cols):
        fig.add_trace(go.Scatter(
            x=df[x_col],
            y=df[y_col],
            mode="lines",
            line=dict(color=colors[i % len(colors)], width=2),
            name=y_col
        ))
    
    fig.update_layout(
        title=title,
        xaxis_title=x_col,
        height=height,
        margin=dict(l=50, r=50, t=50, b=50),
        plot_bgcolor='rgba(0,0,0,0)',
        paper_bgcolor='rgba(0,0,0,0)',
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="right",
            x=1
        )
    )
    
    fig.update_xaxes(showgrid=True, gridwidth=1, gridcolor='lightgray')
    fig.update_yaxes(showgrid=True, gridwidth=1, gridcolor='lightgray')
    
    return fig


def create_bar_chart(
    data: Union[pd.DataFrame, List[Dict]],
    x_col: str,
    y_col: str,
    title: str = "",
    color: str = "#1f77b4",
    height: int = 400,
    orientation: str = "v"
) -> go.Figure:
    """
    创建柱状图
    
    Args:
        data: 数据
        x_col: X轴列名
        y_col: Y轴列名
        title: 图表标题
        color: 柱子颜色
        height: 图表高度
        orientation: 方向 ('v' 垂直, 'h' 水平)
    
    Returns:
        Plotly图表对象
    """
    if isinstance(data, list):
        df = pd.DataFrame(data)
    else:
        df = data
    
    fig = go.Figure()
    
    if orientation == "v":
        fig.add_trace(go.Bar(
            x=df[x_col],
            y=df[y_col],
            marker_color=color,
            name=y_col
        ))
    else:
        fig.add_trace(go.Bar(
            x=df[y_col],
            y=df[x_col],
            orientation='h',
            marker_color=color,
            name=y_col
        ))
    
    fig.update_layout(
        title=title,
        height=height,
        showlegend=False,
        margin=dict(l=50, r=50, t=50, b=50),
        plot_bgcolor='rgba(0,0,0,0)',
        paper_bgcolor='rgba(0,0,0,0)'
    )
    
    fig.update_xaxes(showgrid=True, gridwidth=1, gridcolor='lightgray')
    fig.update_yaxes(showgrid=True, gridwidth=1, gridcolor='lightgray')
    
    return fig


def create_pie_chart(
    data: Union[pd.DataFrame, List[Dict], Dict],
    labels_col: Optional[str] = None,
    values_col: Optional[str] = None,
    title: str = "",
    colors: Optional[List[str]] = None,
    height: int = 400,
    hole: float = 0.0
) -> go.Figure:
    """
    创建饼图
    
    Args:
        data: 数据
        labels_col: 标签列名
        values_col: 数值列名
        title: 图表标题
        colors: 颜色列表
        height: 图表高度
        hole: 中心空洞大小 (0-1)
    
    Returns:
        Plotly图表对象
    """
    if isinstance(data, dict):
        labels = list(data.keys())
        values = list(data.values())
    elif isinstance(data, list):
        df = pd.DataFrame(data)
        labels = df[labels_col].tolist()
        values = df[values_col].tolist()
    else:
        labels = data[labels_col].tolist()
        values = data[values_col].tolist()
    
    fig = go.Figure(data=[go.Pie(
        labels=labels,
        values=values,
        hole=hole,
        marker_colors=colors
    )])
    
    fig.update_layout(
        title=title,
        height=height,
        margin=dict(l=50, r=50, t=50, b=50),
        showlegend=True,
        legend=dict(
            orientation="v",
            yanchor="middle",
            y=0.5,
            xanchor="left",
            x=1.01
        )
    )
    
    return fig


def create_gauge_chart(
    value: float,
    title: str = "",
    min_value: float = 0,
    max_value: float = 100,
    color: str = "#1f77b4",
    height: int = 300
) -> go.Figure:
    """
    创建仪表盘图
    
    Args:
        value: 当前值
        title: 图表标题
        min_value: 最小值
        max_value: 最大值
        color: 颜色
        height: 图表高度
    
    Returns:
        Plotly图表对象
    """
    fig = go.Figure(go.Indicator(
        mode="gauge+number+delta",
        value=value,
        domain={'x': [0, 1], 'y': [0, 1]},
        title={'text': title},
        delta={'reference': max_value * 0.8},
        gauge={
            'axis': {'range': [None, max_value]},
            'bar': {'color': color},
            'steps': [
                {'range': [0, max_value * 0.5], 'color': "lightgray"},
                {'range': [max_value * 0.5, max_value * 0.8], 'color': "yellow"},
                {'range': [max_value * 0.8, max_value], 'color': "lightgreen"}
            ],
            'threshold': {
                'line': {'color': "red", 'width': 4},
                'thickness': 0.75,
                'value': max_value * 0.9
            }
        }
    ))
    
    fig.update_layout(
        height=height,
        margin=dict(l=50, r=50, t=50, b=50),
        paper_bgcolor='rgba(0,0,0,0)'
    )
    
    return fig


def create_heatmap(
    data: Union[pd.DataFrame, np.ndarray],
    title: str = "",
    colorscale: str = "Blues",
    height: int = 400,
    x_labels: Optional[List[str]] = None,
    y_labels: Optional[List[str]] = None
) -> go.Figure:
    """
    创建热力图
    
    Args:
        data: 数据矩阵
        title: 图表标题
        colorscale: 颜色方案
        height: 图表高度
        x_labels: X轴标签
        y_labels: Y轴标签
    
    Returns:
        Plotly图表对象
    """
    if isinstance(data, pd.DataFrame):
        z = data.values
        x_labels = x_labels or data.columns.tolist()
        y_labels = y_labels or data.index.tolist()
    else:
        z = data
    
    fig = go.Figure(data=go.Heatmap(
        z=z,
        x=x_labels,
        y=y_labels,
        colorscale=colorscale,
        showscale=True
    ))
    
    fig.update_layout(
        title=title,
        height=height,
        margin=dict(l=50, r=50, t=50, b=50)
    )
    
    return fig


def create_candlestick_chart(
    data: Union[pd.DataFrame, List[Dict]],
    date_col: str = "date",
    open_col: str = "open",
    high_col: str = "high",
    low_col: str = "low",
    close_col: str = "close",
    volume_col: Optional[str] = None,
    title: str = "",
    height: int = 500
) -> go.Figure:
    """
    创建K线图
    
    Args:
        data: 数据
        date_col: 日期列名
        open_col: 开盘价列名
        high_col: 最高价列名
        low_col: 最低价列名
        close_col: 收盘价列名
        volume_col: 成交量列名
        title: 图表标题
        height: 图表高度
    
    Returns:
        Plotly图表对象
    """
    if isinstance(data, list):
        df = pd.DataFrame(data)
    else:
        df = data
    
    # 创建子图
    if volume_col:
        from plotly.subplots import make_subplots
        fig = make_subplots(
            rows=2, cols=1,
            shared_xaxes=True,
            vertical_spacing=0.03,
            subplot_titles=('价格', '成交量'),
            row_width=[0.2, 0.7]
        )
        
        # K线图
        fig.add_trace(go.Candlestick(
            x=df[date_col],
            open=df[open_col],
            high=df[high_col],
            low=df[low_col],
            close=df[close_col],
            name="价格"
        ), row=1, col=1)
        
        # 成交量图
        fig.add_trace(go.Bar(
            x=df[date_col],
            y=df[volume_col],
            name="成交量",
            marker_color='rgba(158,202,225,0.8)'
        ), row=2, col=1)
        
    else:
        fig = go.Figure(data=[go.Candlestick(
            x=df[date_col],
            open=df[open_col],
            high=df[high_col],
            low=df[low_col],
            close=df[close_col]
        )])
    
    fig.update_layout(
        title=title,
        height=height,
        xaxis_rangeslider_visible=False,
        margin=dict(l=50, r=50, t=50, b=50)
    )
    
    return fig


def create_scatter_plot(
    data: Union[pd.DataFrame, List[Dict]],
    x_col: str,
    y_col: str,
    size_col: Optional[str] = None,
    color_col: Optional[str] = None,
    title: str = "",
    height: int = 400
) -> go.Figure:
    """
    创建散点图
    
    Args:
        data: 数据
        x_col: X轴列名
        y_col: Y轴列名
        size_col: 大小列名
        color_col: 颜色列名
        title: 图表标题
        height: 图表高度
    
    Returns:
        Plotly图表对象
    """
    if isinstance(data, list):
        df = pd.DataFrame(data)
    else:
        df = data
    
    fig = px.scatter(
        df,
        x=x_col,
        y=y_col,
        size=size_col,
        color=color_col,
        title=title,
        height=height
    )
    
    fig.update_layout(
        margin=dict(l=50, r=50, t=50, b=50),
        plot_bgcolor='rgba(0,0,0,0)',
        paper_bgcolor='rgba(0,0,0,0)'
    )
    
    fig.update_xaxes(showgrid=True, gridwidth=1, gridcolor='lightgray')
    fig.update_yaxes(showgrid=True, gridwidth=1, gridcolor='lightgray')
    
    return fig


def create_box_plot(
    data: Union[pd.DataFrame, List[Dict]],
    y_col: str,
    x_col: Optional[str] = None,
    title: str = "",
    height: int = 400
) -> go.Figure:
    """
    创建箱线图
    
    Args:
        data: 数据
        y_col: Y轴列名
        x_col: X轴列名（分组）
        title: 图表标题
        height: 图表高度
    
    Returns:
        Plotly图表对象
    """
    if isinstance(data, list):
        df = pd.DataFrame(data)
    else:
        df = data
    
    fig = px.box(
        df,
        x=x_col,
        y=y_col,
        title=title,
        height=height
    )
    
    fig.update_layout(
        margin=dict(l=50, r=50, t=50, b=50),
        plot_bgcolor='rgba(0,0,0,0)',
        paper_bgcolor='rgba(0,0,0,0)'
    )
    
    return fig


def create_histogram(
    data: Union[pd.DataFrame, List[Dict], List[float]],
    col: Optional[str] = None,
    bins: int = 30,
    title: str = "",
    color: str = "#1f77b4",
    height: int = 400
) -> go.Figure:
    """
    创建直方图
    
    Args:
        data: 数据
        col: 列名
        bins: 分箱数量
        title: 图表标题
        color: 颜色
        height: 图表高度
    
    Returns:
        Plotly图表对象
    """
    if isinstance(data, list) and col is None:
        values = data
    elif isinstance(data, list):
        df = pd.DataFrame(data)
        values = df[col].tolist()
    else:
        values = data[col].tolist()
    
    fig = go.Figure(data=[go.Histogram(
        x=values,
        nbinsx=bins,
        marker_color=color
    )])
    
    fig.update_layout(
        title=title,
        height=height,
        margin=dict(l=50, r=50, t=50, b=50),
        plot_bgcolor='rgba(0,0,0,0)',
        paper_bgcolor='rgba(0,0,0,0)'
    )
    
    fig.update_xaxes(showgrid=True, gridwidth=1, gridcolor='lightgray')
    fig.update_yaxes(showgrid=True, gridwidth=1, gridcolor='lightgray')
    
    return fig


def create_waterfall_chart(
    categories: List[str],
    values: List[float],
    title: str = "",
    height: int = 400
) -> go.Figure:
    """
    创建瀑布图
    
    Args:
        categories: 分类列表
        values: 数值列表
        title: 图表标题
        height: 图表高度
    
    Returns:
        Plotly图表对象
    """
    fig = go.Figure(go.Waterfall(
        name="",
        orientation="v",
        measure=["relative"] * (len(categories) - 1) + ["total"],
        x=categories,
        textposition="outside",
        text=[f"{v:+.1f}" for v in values],
        y=values,
        connector={"line": {"color": "rgb(63, 63, 63)"}},
    ))
    
    fig.update_layout(
        title=title,
        height=height,
        margin=dict(l=50, r=50, t=50, b=50),
        showlegend=False
    )
    
    return fig